import os
from collections import defaultdict
from PIL import Image
from . import DATASET
from lib.base import BaseDataset
import torch
import sys
import numpy as np
from torchvision.transforms import transforms


@DATASET.register('cub200')
class CUBDataset(BaseDataset):

    def __init__(self, cfg, type):
        # self.part_size = (84, 84)

        super().__init__(cfg, type)

    def _get_split_categories(self) -> dict:
        categories = list(self.category2path)
        category_split = defaultdict(list)
        for i, cate in enumerate(categories):
            if i % 2 == 0:
                category_split['TRAIN'].append(cate)
            if i % 4 == 1:
                category_split['VALIDATE'].append(cate)
            if i % 4 == 3:
                category_split['TEST'].append(cate)
        return category_split

    def _get_category2path(self):
        labels = os.listdir(os.path.join(self.data_path, 'images'))
        labels.sort()

        category2path = defaultdict(list)
        id2category = {i: labels[i] for i in range(len(labels))}

        with open(os.path.join(self.data_path, 'image_class_labels.txt')) as f1:
            with open(os.path.join(self.data_path, 'images.txt')) as f2:
                for line1, line2 in zip(f1, f2):
                    _, id = line1.split()
                    _, path = line2.split()
                    category2path[id2category[int(id) - 1]].append(
                        os.path.join(self.data_path, 'images', path))
        return category2path
